# -*- coding: utf-8 -*-

import os
import shutil
import zipfile

# from Composition_code import do_composite
from Composition_code_test import do_composite_test

if __name__ == '__main__':
    # path to provided foreground images
    fg_path = 'data/fg/'
    # path to provided alpha mattes
    a_path = 'data/mask/'
    # Path to background images (MSCOCO)
    bg_path = 'data/bg/'
    # Path to folder where you want the composited images to go
    out_path = 'data/merged/'

    train_folder = 'data/'

    if not os.path.exists(bg_path):
        zip_file = 'data/train2014.zip'
        print('Extracting {}...'.format(zip_file))

        zip_ref = zipfile.ZipFile(zip_file, 'r')
        zip_ref.extractall('data')
        zip_ref.close()

        with open(os.path.join(train_folder, 'training_bg_names.txt')) as f:
            training_bg_names = f.read().splitlines()

        os.makedirs(bg_path)
        for bg_name in training_bg_names:
            src_path = os.path.join('data/train2014', bg_name)
            dest_path = os.path.join(bg_path, bg_name)
            shutil.move(src_path, dest_path)

    if not os.path.exists(out_path):
        os.makedirs(out_path)
    # do_composite()

    # path to provided foreground images
    fg_test_path = 'data/fg_test/'
    # path to provided alpha mattes
    a_test_path = 'data/mask_test/'
    # Path to background images (PASCAL VOC)
    bg_test_path = 'data/bg_test/'
    # Path to folder where you want the composited images to go
    out_test_path = 'data/merged_test/'

    # test data gen
    test_folder = 'data/'

    if not os.path.exists(bg_test_path):
        os.makedirs(bg_test_path)

    with open(os.path.join(test_folder, 'test_bg_names.txt')) as f:
        test_bg_names = f.read().splitlines()

    if not os.path.exists(out_test_path):
        os.makedirs(out_test_path)

    do_composite_test()
